import io
from typing import List

import cv2
import numpy as np
from .preprocessing import Preprocessing

class VideoPreprocessor(Preprocessing):
    def __init__(self, target_size=(640, 640), batch_size=20):
        self.target_size = target_size
        self.batch_size = batch_size

    def preprocess(self, video_bytes: bytes) -> np.ndarray:
        """
        预处理视频（以MP4或其他类似格式编码），为Triton推断做准备。

        参数:
            video_bytes: 编码后的视频数据。

        返回值:
            np.ndarray: 预处理后的视频帧，以numpy数组形式返回，数组形状为(num_frames, height, width, channels)。
        """
        video = cv2.VideoCapture(io.BytesIO(video_bytes))
        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        processed_frames = []

        if total_frames < 18000:
            num_batches = -(-total_frames // self.batch_size)  # 向上取整除法
        else:
            num_batches = 0
            self.batch_size = total_frames

        for i in range(num_batches):
            start_index = i * self.batch_size
            end_index = min((i + 1) * self.batch_size, total_frames)

            batch_frames = self._process_batch(video, start_index, end_index)
            processed_frames.extend(batch_frames)

        video.release()

        return np.concatenate(processed_frames, axis=0)

    def _process_batch(self, video: cv2.VideoCapture, start_index: int, end_index: int) -> List[np.ndarray]:
        video.set(cv2.CAP_PROP_POS_FRAMES, start_index)
        frames = []

        for _ in range(start_index, end_index):
            ret, frame = video.read()
            if not ret:
                break

            resized_frame = cv2.resize(frame, self.target_size)
            frames.append(resized_frame)

        return frames
